package com.bleeth.flow.step.row;


import com.bleeth.flow.core.common.APlugin;
import com.bleeth.flow.core.common.PluginAnnotation;
import com.bleeth.flow.core.common.PluginTypeEnum;
import com.jd.platform.async.worker.WorkResult;
import com.jd.platform.async.wrapper.WorkerWrapper;
import lombok.Data;
import org.apache.spark.sql.Column;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;

import java.io.Serializable;
import java.util.List;
import java.util.Map;

/**
 * @author ：Bleeth
 * @date ：2021-08-05
 * @description：数据关联组件
 */
@Data
@PluginAnnotation(name = "数据关联组件",
        type = PluginTypeEnum.ROW,
        description = "",
        id = "JoinPlugin")
public class JoinPlugin extends APlugin implements Serializable {

    private static final long serialVersionUID = -4488664272322183836L;

    private List<JoinParam> joinParamList;

    private String leftPlugin;

    private String rightPlugin;


    //默认为 `inner`, `cross`, `outer`, `full`, `full_outer`, `left`, `left_outer`,
    //`right`, `right_outer`, `left_semi`, `left_anti`.
    private String joinType;




    @Override
    public Dataset<Row> action(Dataset<Row> input, Map<String, WorkerWrapper> allWrappers) {
        WorkerWrapper leftWrapper = allWrappers.get(leftPlugin);
        WorkerWrapper rightWrapper = allWrappers.get(rightPlugin);

        WorkResult leftWorkResult = leftWrapper.getWorkResult();
        WorkResult rightWorkResult = rightWrapper.getWorkResult();

        Dataset<Row> leftDs = (Dataset<Row>) leftWorkResult.getResult();
        Dataset<Row> rightDs = (Dataset<Row>) rightWorkResult.getResult();

        //join组件实现left join,right join,full join
        Column joinExpr = null;
        for (JoinParam joinParam : joinParamList) {
            String leftField = joinParam.getLeftField();
            String rightField = joinParam.getRightField();
            Column otherExpr = leftDs.apply(leftField).equalTo(rightDs.apply(rightField));
            if (joinExpr == null) {
                joinExpr = otherExpr;
                continue;
            }
            joinExpr.and(otherExpr);
        }

        Dataset<Row> ds = leftDs.join(rightDs, joinExpr, joinType);
        return ds;
    }


    @Data
    public static class JoinParam implements Serializable {
        private static final long serialVersionUID = -7590680435186574933L;

        private String leftField;

        private String rightField;

    }


}
